# ques4_prepare_data/sheet2_女胎检测数据_fillnan_sifted_outlier.csv
import warnings
warnings.filterwarnings('ignore')
import util_for_output_zh

import os,pdb
import pandas as pd
import numpy as np
# import matplotlib.pyplot as plt
from util_set_zh_matplot import plt
import seaborn as sns
from sklearn.linear_model import Ridge  # 改用更简单的模型
from sklearn.preprocessing import PolynomialFeatures, StandardScaler
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_absolute_error, mean_squared_error
from skopt import gp_minimize
from skopt.space import Real
from typing import Tuple, List, Dict
from joblib import Parallel, delayed  # 并行处理
from sklearn.cluster import KMeans  # 用于预分组

# 设置显示选项
pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', 100)
pd.set_option('display.width', None)
pd.set_option('display.max_colwidth', None)

# 创建输出目录
out_dir = 'ques4_prepare_data'
os.makedirs(out_dir, exist_ok=True)


from sklearn.ensemble import RandomForestClassifier
from sklearn.preprocessing import LabelEncoder, StandardScaler
from sklearn.metrics import classification_report, accuracy_score
import xgboost as xgb

# 1. 数据加载与预处理
def load_and_preprocess(filepath):
    # 读取数据
    data = pd.read_csv(filepath)
    
    # 选择需要的特征和目标列
    features = [
        '原始读段数', '在参考基因组上比对的比例', '重复读段的比例', '唯一比对的读段数',
        'GC含量', '13号染色体的Z值', '18号染色体的Z值', '21号染色体的Z值', 
        'X染色体的Z值', 'X染色体浓度', '13号染色体的GC含量', 
        '18号染色体的GC含量', '21号染色体的GC含量', '孕妇BMI'
    ]
    target = '染色体的非整倍体'
    
    X = data[features]
    y = data[target]
    
    # 编码目标变量
    le = LabelEncoder()
    y_encoded = le.fit_transform(y)
    
    # 标准化特征
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    
    return X_scaled, y_encoded, X, y, le, scaler

# 2. 特征重要性分析
def analyze_feature_importance(X, y, feature_names):
    # 使用随机森林分析特征重要性
    rf = RandomForestClassifier(n_estimators=100, random_state=42)
    rf.fit(X, y)
    
    # 获取特征重要性
    importances = rf.feature_importances_
    indices = np.argsort(importances)[::-1]
    
    # 打印特征重要性
    print("特征重要性排序:")
    for f in range(X.shape[1]):
        print(f"{f+1}. {feature_names[indices[f]]}: {importances[indices[f]]:.4f}")
    
    # 可视化特征重要性
    plt.figure(figsize=(10, 6))
    plt.title("特征重要性")
    sns.barplot(x=importances[indices], y=np.array(feature_names)[indices])
    plt.tight_layout()
    # plt.show()
    plt.savefig( os.path.join( out_dir , 'ques4_特征重要性.png') )
    
    return rf

# 3. 构建并评估模型
def build_and_evaluate_model(X, y):
    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    
    # 随机森林模型
    rf_model = RandomForestClassifier(n_estimators=200, random_state=42)
    rf_model.fit(X_train, y_train)
    
    # XGBoost模型
    xgb_model = xgb.XGBClassifier(objective='multi:softmax', num_class=len(np.unique(y)), random_state=42)
    xgb_model.fit(X_train, y_train)
    
    # 评估模型
    for name, model in [('随机森林', rf_model), ('XGBoost', xgb_model)]:
        y_pred = model.predict(X_test)
        print(f"\n{name}模型性能:")
        print(f"准确率: {accuracy_score(y_test, y_pred):.4f}")
        print(classification_report(y_test, y_pred))
    
    return rf_model, xgb_model

# 4. 染色体异常判定方法
class ChromosomeAbnormalityDetector:
    def __init__(self, model, scaler, label_encoder):
        self.model = model
        self.scaler = scaler
        self.label_encoder = label_encoder
    
    def predict(self, features):
        # 输入特征顺序应与训练时一致
        feature_names = [
            '原始读段数', '在参考基因组上比对的比例', '重复读段的比例', '唯一比对的读段数',
            'GC含量', '13号染色体的Z值', '18号染色体的Z值', '21号染色体的Z值', 
            'X染色体的Z值', 'X染色体浓度', '13号染色体的GC含量', 
            '18号染色体的GC含量', '21号染色体的GC含量', '孕妇BMI'
        ]
        
        # 将输入转换为DataFrame确保顺序正确
        input_df = pd.DataFrame([features], columns=feature_names)
        
        # 标准化特征
        scaled_features = self.scaler.transform(input_df)
        
        # 预测
        prediction = self.model.predict(scaled_features)
        proba = self.model.predict_proba(scaled_features)
        
        # 解码预测结果
        predicted_label = self.label_encoder.inverse_transform(prediction)[0]
        confidence = np.max(proba)
        
        return predicted_label, confidence
    
# 主程序
if __name__ == "__main__":
    # 文件路径 - 请替换为实际路径
    filepath = 'ques4_prepare_data/sheet2_女胎检测数据_fillnan_sifted_outlier.csv'
    
    # 1. 数据预处理
    X_scaled, y_encoded, X, y, le, scaler = load_and_preprocess(filepath)
    
    # 2. 特征重要性分析
    feature_names = X.columns.tolist()
    rf_model = analyze_feature_importance(X_scaled, y_encoded, feature_names)
    
    # 3. 构建并评估模型
    rf_model, xgb_model = build_and_evaluate_model(X_scaled, y_encoded)
    
    # 4. 创建异常检测器实例
    detector = ChromosomeAbnormalityDetector(rf_model, scaler, le)
    
    # 示例预测
    sample_features = [
        5000000, 0.81, 0.03, 4000000,  # 测序相关特征
        0.40, 0.9, -0.1, 0.8, 1.2,     # Z值相关特征
        -0.02, 0.38, 0.39, 0.40,       # GC含量相关特征
        30                              # BMI
    ]
    
    prediction, confidence = detector.predict(sample_features)
    print(f"\n示例预测结果: {prediction} (置信度: {confidence:.2f})")
    
    # 保存模型供后续使用
    import joblib
    joblib.dump(detector, os.path.join(out_dir , 'chromosome_abnormality_detector.pkl'))